import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'
import tensorflow as tf
from tensorflow.python import keras

"""
第一层
卷积：[None, 32, 32, 3]———>[None, 32, 32, 32]
权重数量：[5, 5, 1 ,32]
偏置数量：[32]
激活：[None, 32, 32, 32]———>[None, 32, 32, 32]
池化：[None, 32, 32, 32]———>[None, 16, 16, 32]
第二层
卷积：[None, 16, 16, 32]———>[None, 16, 16, 64]
权重数量：[5, 5, 32 ,64]
偏置数量：[64]
激活：[None, 16, 16, 64]———>[None, 16, 16, 64]
池化：[None, 16, 16, 64]———>[None, 8, 8, 64]
全连接层
[None, 8, 8, 64]——>[None, 8 8 64]
[None, 8 8 64] x [8 8 64, 1024] = [None, 1024]
[None,1024] x [1024, 100]——>[None, 100]
权重数量：[8 8 64, 1024] + [1024, 100]，由分类别数而定
偏置数量：[1024] + [100]，由分类别数而定
"""


class CNNMinist:

    # 构建模型
    model = keras.Sequential([
        # 卷积一层：32个5*5*3的filter，步长为1，padding为same
        #
        keras.layers.Conv2D(filters=32, kernel_size=5, strides=1, padding="same", data_format="channels_last",
                            activation=tf.nn.relu,
                            # kernelregularizer=keras.regularizers.l2()  # # 添加L2正则化
                            ),
        # 池化一层：
        keras.layers.MaxPool2D(pool_size=2, strides=2, padding="same"),
        # 卷积二层：64个5*5*32的filter，步长为1， padding为same
        keras.layers.Conv2D(filters=64, kernel_size=5, strides=1, padding="same", data_format="channels_last",
                            activation=tf.nn.relu),
        # 池化二层：
        keras.layers.MaxPool2D(pool_size=2, strides=2, padding="same"),

        # 全连接层盛景网络
        # 将数据变换为二维形状
        keras.layers.Flatten(),
        # 1024个神经元的网络层
        keras.layers.Dense(1024, activation=tf.nn.relu),
        # 100个分类的全连接层
        keras.layers.Dense(100, activation=tf.nn.softmax)
    ])

    def __init__(self):
        # 获取数据
        (self.x_train, self.y_train), (self.x_test, self.y_test) = keras.datasets.cifar100.load_data()

        # 数据归一化
        self.x_train = self.x_train.reshape((-1, 32, 32, 3)) / 255.0
        self.x_test = self.x_test.reshape((-1, 32, 32, 3)) / 255.0

    def compile(self):
        """
        编译模型创建优化器，损失计算，准确率
        :return:
        """
        CNNMinist.model.compile(optimizer=keras.optimizers.Adam(),
                                loss=keras.losses.sparse_categorical_crossentropy,
                                metrics=["accuracy"])

    def fit(self):
        """
        模型训练
        :return:
        """
        CNNMinist.model.fit(self.x_train, self.y_train, batch_size=64, epochs=1)

        CNNMinist.model.save_weights("./test_ckpt/cnn.h5")

    def evaluate(self):
        """
        评估模型
        :return:
        """
        # model = CNNMinist.model.load_weights("./test_ckpt/cnn.h5")
        # print(model)
        loss, acc = CNNMinist.model.evaluate(self.x_test, self.y_test, )
        print("损失为:", loss)
        print("准确率为:", acc)


if __name__ == '__main__':
    cnn = CNNMinist()
    cnn.compile()
    cnn.fit()
    cnn.evaluate()
    print(CNNMinist.model.summary())


